{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 编译 PyTorch 模型\n",
    "\n",
    "**Author**: [Yaoda Zhou](https://github.com/juda)\n",
    "\n",
    "本文是一篇使用装饰器`optimize_torch`优化PyTorch模型的教程。要跟随本教程，需要安装 PyTorch 以及 TorchVision：\n",
    "```bash\n",
    "%%shell\n",
    "pip install torch\n",
    "pip install torchvision\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm/python/tvm/contrib/torch/__init__.py:50: RuntimeWarning: The library libpt_tvmdsoop is not built successfully. /media/pc/data/lxw/ai/tvm/build/libpt_tvmdsoop.so: cannot open shared object file: No such file or directory\n",
      "  warnings.warn(\n",
      "/media/pc/data/lxw/ai/tvm/python/tvm/contrib/torch/__init__.py:50: RuntimeWarning: The library libpt_tvmdsoop_new is not built successfully. /media/pc/data/lxw/ai/tvm/build/libpt_tvmdsoop_new.so: cannot open shared object file: No such file or directory\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# Import library for profiling\n",
    "import torch.utils.benchmark as benchmark\n",
    "from torchvision.models import resnet18\n",
    "\n",
    "# Import `optimize_torch` function\n",
    "from tvm.contrib.torch import optimize_torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 PyTorch 构建简单模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 20, 5)\n",
    "        self.conv2 = nn.Conv2d(20, 20, 5)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.conv1(x))\n",
    "        return F.relu(self.conv2(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 TVM MetaSchedule 优化 SimpleModel\n",
    "\n",
    "我们提供了`optimize_torch`函数，其用法与`torch.jit.trace`类似。用户需要提供要优化的PyTorch模型以及其示例输入。PyTorch模块将由TVM针对目标硬件进行调优。如果不提供额外信息，模型将针对CPU进行调优。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-20 12:21:19 [INFO] Logging directory: /tmp/tmpl0j3jqte/logs\n",
      "2024-03-20 12:21:36 [INFO] LocalBuilder: max_workers = 24\n",
      "2024-03-20 12:21:38 [INFO] LocalRunner: max_workers = 1\n",
      "2024-03-20 12:21:39 [INFO] [task_scheduler.cc:159] Initializing Task #0: \"fused_layout_transform\"\n",
      "2024-03-20 12:21:39 [INFO] [task_scheduler.cc:159] Initializing Task #1: \"fused_nn_contrib_conv2d_NCHWc_add_nn_relu\"\n",
      "2024-03-20 12:21:39 [INFO] [task_scheduler.cc:159] Initializing Task #2: \"fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1\"\n",
      "2024-03-20 12:21:39 [INFO] [task_scheduler.cc:159] Initializing Task #3: \"fused_layout_transform_1\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fused_layout_transform</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu</td>\n",
       "      <td>748800</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1</td>\n",
       "      <td>1603200</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>fused_layout_transform_1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           Name       FLOP    Weight   \\\n",
       "0                        fused_layout_transform          1         1    \n",
       "1     fused_nn_contrib_conv2d_NCHWc_add_nn_relu     748800         1    \n",
       "2   fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1    1603200         1    \n",
       "3                      fused_layout_transform_1          1         1    \n",
       "\n",
       "    Speed (GFLOPS)    Latency (us)    Weighted Latency (us)    Trials    Done   \n",
       "0              N/A             N/A                      N/A         0           \n",
       "1              N/A             N/A                      N/A         0           \n",
       "2              N/A             N/A                      N/A         0           \n",
       "3              N/A             N/A                      N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-20 12:21:40 [DEBUG] [task_scheduler.cc:318] \n",
      " ID |                                        Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "  0 |                      fused_layout_transform |       1 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  1 |   fused_nn_contrib_conv2d_NCHWc_add_nn_relu |  748800 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  2 | fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1 | 1603200 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  3 |                    fused_layout_transform_1 |       1 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "\n",
      "Total trials: 0\n",
      "Total latency (us): 0\n",
      "\n",
      "2024-03-20 12:21:40 [INFO] [task_scheduler.cc:180] TaskScheduler picks Task #0: \"fused_layout_transform\"\n",
      "2024-03-20 12:21:40 [INFO] [task_scheduler.cc:193] Sending 2 sample(s) to builder\n",
      "2024-03-20 12:21:42 [INFO] [task_scheduler.cc:195] Sending 2 sample(s) to runner\n",
      "2024-03-20 12:21:43 [DEBUG] XGB iter   0: tr-p-rmse: 0.424805\ttr-a-peak@32: 1.000000\ttr-rmse: 0.424910\ttr-rmse: 0.424910\n",
      "2024-03-20 12:21:43 [DEBUG] XGB iter  25: tr-p-rmse: 0.015707\ttr-a-peak@32: 1.000000\ttr-rmse: 0.015787\ttr-rmse: 0.015787\n",
      "2024-03-20 12:21:43 [DEBUG] XGB iter  50: tr-p-rmse: 0.010228\ttr-a-peak@32: 1.000000\ttr-rmse: 0.010230\ttr-rmse: 0.010230\n",
      "2024-03-20 12:21:43 [DEBUG] XGB iter  75: tr-p-rmse: 0.010225\ttr-a-peak@32: 1.000000\ttr-rmse: 0.010224\ttr-rmse: 0.010224\n",
      "2024-03-20 12:21:43 [DEBUG] XGB iter 100: tr-p-rmse: 0.010226\ttr-a-peak@32: 1.000000\ttr-rmse: 0.010224\ttr-rmse: 0.010224\n",
      "2024-03-20 12:21:43 [DEBUG] XGB stopped. Best iteration: [54] tr-p-rmse:0.01022\ttr-a-peak@32:1.00000\ttr-rmse:0.01023\ttr-rmse:0.01023 \n",
      "2024-03-20 12:21:43 [INFO] [task_scheduler.cc:237] [Updated] Task #0: \"fused_layout_transform\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fused_layout_transform</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>2</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu</td>\n",
       "      <td>748800</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1</td>\n",
       "      <td>1603200</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>fused_layout_transform_1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           Name       FLOP    Weight   \\\n",
       "0                        fused_layout_transform          1         1    \n",
       "1     fused_nn_contrib_conv2d_NCHWc_add_nn_relu     748800         1    \n",
       "2   fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1    1603200         1    \n",
       "3                      fused_layout_transform_1          1         1    \n",
       "\n",
       "    Speed (GFLOPS)    Latency (us)    Weighted Latency (us)    Trials    Done   \n",
       "0           0.0001         11.0077                  11.0077         2           \n",
       "1              N/A             N/A                      N/A         0           \n",
       "2              N/A             N/A                      N/A         0           \n",
       "3              N/A             N/A                      N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-20 12:21:43 [DEBUG] [task_scheduler.cc:318] \n",
      " ID |                                        Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "  0 |                      fused_layout_transform |       1 |      1 |         0.0001 |      11.0077 |               11.0077 |      2 |      \n",
      "  1 |   fused_nn_contrib_conv2d_NCHWc_add_nn_relu |  748800 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  2 | fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1 | 1603200 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  3 |                    fused_layout_transform_1 |       1 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "2024-03-20 12:21:43 [INFO] [task_scheduler.cc:260] Task #0 has finished. Remaining task(s): 3\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fused_layout_transform</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>2</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu</td>\n",
       "      <td>748800</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1</td>\n",
       "      <td>1603200</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>fused_layout_transform_1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           Name       FLOP    Weight   \\\n",
       "0                        fused_layout_transform          1         1    \n",
       "1     fused_nn_contrib_conv2d_NCHWc_add_nn_relu     748800         1    \n",
       "2   fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1    1603200         1    \n",
       "3                      fused_layout_transform_1          1         1    \n",
       "\n",
       "    Speed (GFLOPS)    Latency (us)    Weighted Latency (us)    Trials    Done   \n",
       "0           0.0001         11.0077                  11.0077         2       Y   \n",
       "1              N/A             N/A                      N/A         0           \n",
       "2              N/A             N/A                      N/A         0           \n",
       "3              N/A             N/A                      N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-20 12:21:43 [DEBUG] [task_scheduler.cc:318] \n",
      " ID |                                        Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "  0 |                      fused_layout_transform |       1 |      1 |         0.0001 |      11.0077 |               11.0077 |      2 |    Y \n",
      "  1 |   fused_nn_contrib_conv2d_NCHWc_add_nn_relu |  748800 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  2 | fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1 | 1603200 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  3 |                    fused_layout_transform_1 |       1 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "2024-03-20 12:21:43 [INFO] [task_scheduler.cc:260] Task #1 has finished. Remaining task(s): 2\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fused_layout_transform</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>2</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu</td>\n",
       "      <td>748800</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1</td>\n",
       "      <td>1603200</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>fused_layout_transform_1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           Name       FLOP    Weight   \\\n",
       "0                        fused_layout_transform          1         1    \n",
       "1     fused_nn_contrib_conv2d_NCHWc_add_nn_relu     748800         1    \n",
       "2   fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1    1603200         1    \n",
       "3                      fused_layout_transform_1          1         1    \n",
       "\n",
       "    Speed (GFLOPS)    Latency (us)    Weighted Latency (us)    Trials    Done   \n",
       "0           0.0001         11.0077                  11.0077         2       Y   \n",
       "1              N/A             N/A                      N/A         0       Y   \n",
       "2              N/A             N/A                      N/A         0           \n",
       "3              N/A             N/A                      N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-20 12:21:43 [DEBUG] [task_scheduler.cc:318] \n",
      " ID |                                        Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "  0 |                      fused_layout_transform |       1 |      1 |         0.0001 |      11.0077 |               11.0077 |      2 |    Y \n",
      "  1 |   fused_nn_contrib_conv2d_NCHWc_add_nn_relu |  748800 |      1 |            N/A |          N/A |                   N/A |      0 |    Y \n",
      "  2 | fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1 | 1603200 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "  3 |                    fused_layout_transform_1 |       1 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "2024-03-20 12:21:43 [INFO] [task_scheduler.cc:260] Task #2 has finished. Remaining task(s): 1\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fused_layout_transform</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>2</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu</td>\n",
       "      <td>748800</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1</td>\n",
       "      <td>1603200</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>fused_layout_transform_1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           Name       FLOP    Weight   \\\n",
       "0                        fused_layout_transform          1         1    \n",
       "1     fused_nn_contrib_conv2d_NCHWc_add_nn_relu     748800         1    \n",
       "2   fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1    1603200         1    \n",
       "3                      fused_layout_transform_1          1         1    \n",
       "\n",
       "    Speed (GFLOPS)    Latency (us)    Weighted Latency (us)    Trials    Done   \n",
       "0           0.0001         11.0077                  11.0077         2       Y   \n",
       "1              N/A             N/A                      N/A         0       Y   \n",
       "2              N/A             N/A                      N/A         0       Y   \n",
       "3              N/A             N/A                      N/A         0           "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-20 12:21:43 [DEBUG] [task_scheduler.cc:318] \n",
      " ID |                                        Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "  0 |                      fused_layout_transform |       1 |      1 |         0.0001 |      11.0077 |               11.0077 |      2 |    Y \n",
      "  1 |   fused_nn_contrib_conv2d_NCHWc_add_nn_relu |  748800 |      1 |            N/A |          N/A |                   N/A |      0 |    Y \n",
      "  2 | fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1 | 1603200 |      1 |            N/A |          N/A |                   N/A |      0 |    Y \n",
      "  3 |                    fused_layout_transform_1 |       1 |      1 |            N/A |          N/A |                   N/A |      0 |      \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "2024-03-20 12:21:43 [INFO] [task_scheduler.cc:260] Task #3 has finished. Remaining task(s): 0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>FLOP</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Speed (GFLOPS)</th>\n",
       "      <th>Latency (us)</th>\n",
       "      <th>Weighted Latency (us)</th>\n",
       "      <th>Trials</th>\n",
       "      <th>Done</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fused_layout_transform</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>11.0077</td>\n",
       "      <td>2</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu</td>\n",
       "      <td>748800</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1</td>\n",
       "      <td>1603200</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>fused_layout_transform_1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>N/A</td>\n",
       "      <td>0</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           Name       FLOP    Weight   \\\n",
       "0                        fused_layout_transform          1         1    \n",
       "1     fused_nn_contrib_conv2d_NCHWc_add_nn_relu     748800         1    \n",
       "2   fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1    1603200         1    \n",
       "3                      fused_layout_transform_1          1         1    \n",
       "\n",
       "    Speed (GFLOPS)    Latency (us)    Weighted Latency (us)    Trials    Done   \n",
       "0           0.0001         11.0077                  11.0077         2       Y   \n",
       "1              N/A             N/A                      N/A         0       Y   \n",
       "2              N/A             N/A                      N/A         0       Y   \n",
       "3              N/A             N/A                      N/A         0       Y   "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-03-20 12:21:43 [DEBUG] [task_scheduler.cc:318] \n",
      " ID |                                        Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Done \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "  0 |                      fused_layout_transform |       1 |      1 |         0.0001 |      11.0077 |               11.0077 |      2 |    Y \n",
      "  1 |   fused_nn_contrib_conv2d_NCHWc_add_nn_relu |  748800 |      1 |            N/A |          N/A |                   N/A |      0 |    Y \n",
      "  2 | fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1 | 1603200 |      1 |            N/A |          N/A |                   N/A |      0 |    Y \n",
      "  3 |                    fused_layout_transform_1 |       1 |      1 |            N/A |          N/A |                   N/A |      0 |    Y \n",
      "---------------------------------------------------------------------------------------------------------------------------------------------\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n",
      "\n",
      "Total trials: 2\n",
      "Total latency (us): 11.0077\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[12:21:44] /media/pc/data/lxw/ai/tvm/src/relay/backend/te_compiler_cache.cc:679: Warning: Cannot find workload: fused_nn_contrib_conv2d_NCHWc_add_nn_relu\n",
      "[12:21:44] /media/pc/data/lxw/ai/tvm/src/relay/backend/te_compiler_cache.cc:679: Warning: Cannot find workload: fused_nn_contrib_conv2d_NCHWc_add_nn_relu\n",
      "[12:21:44] /media/pc/data/lxw/ai/tvm/src/relay/backend/te_compiler_cache.cc:679: Warning: Cannot find workload: fused_layout_transform\n",
      "[12:21:44] /media/pc/data/lxw/ai/tvm/src/relay/backend/te_compiler_cache.cc:679: Warning: Cannot find workload: tvmgen_default_fused_nn_contrib_conv2d_NCHWc_add_nn_relu\n",
      "[12:21:44] /media/pc/data/lxw/ai/tvm/src/relay/backend/te_compiler_cache.cc:679: Warning: Cannot find workload: tvmgen_default_fused_nn_contrib_conv2d_NCHWc_add_nn_relu_1\n",
      "[12:21:44] /media/pc/data/lxw/ai/tvm/src/relay/backend/te_compiler_cache.cc:679: Warning: Cannot find workload: tvmgen_default_fused_layout_transform_1\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "optimize_torch requires the flag /\"USE_PT_TVMDSOOP/\" set in config.cmake",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[6], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m simple_model \u001b[38;5;241m=\u001b[39m SimpleModel()\n\u001b[1;32m      2\u001b[0m example_input \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrandn(\u001b[38;5;241m20\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m10\u001b[39m, \u001b[38;5;241m10\u001b[39m)\n\u001b[0;32m----> 3\u001b[0m model_optimized_by_tvm \u001b[38;5;241m=\u001b[39m \u001b[43moptimize_torch\u001b[49m\u001b[43m(\u001b[49m\u001b[43msimple_model\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexample_input\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_trials_global\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/contrib/torch/optimize_torch.py:166\u001b[0m, in \u001b[0;36moptimize_torch\u001b[0;34m(func, example_inputs, max_trials_global, work_dir, target, max_trials_per_task, num_trials_per_iter, builder, runner, database, cost_model, measure_callbacks, task_scheduler, space, strategy, seed)\u001b[0m\n\u001b[1;32m    164\u001b[0m save_runtime_mod \u001b[38;5;241m=\u001b[39m get_global_func(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtvmtorch.save_runtime_mod\u001b[39m\u001b[38;5;124m\"\u001b[39m, allow_missing\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m save_runtime_mod \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 166\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124moptimize_torch requires the flag /\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUSE_PT_TVMDSOOP/\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m set in config.cmake\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m    167\u001b[0m save_runtime_mod(executor_factory\u001b[38;5;241m.\u001b[39mmodule)\n\u001b[1;32m    169\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m GraphExecutorFactoryWrapper(torch\u001b[38;5;241m.\u001b[39mclasses\u001b[38;5;241m.\u001b[39mtvm_torch\u001b[38;5;241m.\u001b[39mGraphExecutorFactoryWrapper())\n",
      "\u001b[0;31mValueError\u001b[0m: optimize_torch requires the flag /\"USE_PT_TVMDSOOP/\" set in config.cmake"
     ]
    }
   ],
   "source": [
    "simple_model = SimpleModel()\n",
    "example_input = torch.randn(20, 1, 10, 10)\n",
    "model_optimized_by_tvm = optimize_torch(simple_model, example_input, max_trials_global=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 保存/加载模块\n",
    "\n",
    "我们可以像标准的`nn.Module`一样保存和加载我们优化过的模块。\n",
    "\n",
    "让我们运行我们的优化模块。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'model_optimized_by_tvm' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m ret1 \u001b[38;5;241m=\u001b[39m \u001b[43mmodel_optimized_by_tvm\u001b[49m(example_input)\n\u001b[1;32m      3\u001b[0m torch\u001b[38;5;241m.\u001b[39msave(model_optimized_by_tvm, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_optimized.pt\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      4\u001b[0m model_loaded \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mload(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_optimized.pt\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'model_optimized_by_tvm' is not defined"
     ]
    }
   ],
   "source": [
    "ret1 = model_optimized_by_tvm(example_input)\n",
    "\n",
    "torch.save(model_optimized_by_tvm, \"model_optimized.pt\")\n",
    "model_loaded = torch.load(\"model_optimized.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# We load the module and run it again.\n",
    "ret2 = model_loaded(example_input)\n",
    "\n",
    "# We will show 2 results:\n",
    "# (1) we can safely load and save model by showing the result of model\n",
    "# after save and load operations is still the same as original one;\n",
    "# (2) the model we optimize returns the same result as the original PyTorch model.\n",
    "\n",
    "ret3 = simple_model(example_input)\n",
    "testing.assert_allclose(ret1.detach().numpy(), ret2.detach().numpy(), atol=1e-5, rtol=1e-5)\n",
    "testing.assert_allclose(ret1.detach().numpy(), ret3.detach().numpy(), atol=1e-5, rtol=1e-5)\n",
    "\n",
    "######################################################################\n",
    "# Optimize resnet18\n",
    "# -----------------\n",
    "# In the following, we will show that our approach is able to\n",
    "# accelerate common models, such as resnet18.\n",
    "\n",
    "# We will tune our model for the GPU.\n",
    "target_cuda = \"nvidia/geforce-rtx-3070\"\n",
    "\n",
    "# For PyTorch users, the code could be written as usual, except for\n",
    "# applying \"optimize_torch\" function on the resnet18 model.\n",
    "\n",
    "resnet18_tvm = optimize_torch(\n",
    "    resnet18().cuda().eval(), [torch.rand(1, 3, 224, 224).cuda()], target=target_cuda\n",
    ")\n",
    "\n",
    "# TorchScript also provides a built-in \"optimize_for_inference\" function to accelerate the inference.\n",
    "resnet18_torch = torch.jit.optimize_for_inference(torch.jit.script(resnet18().cuda().eval()))\n",
    "\n",
    "\n",
    "######################################################################\n",
    "# Compare the performance between two approaches\n",
    "# ----------------------------------------------\n",
    "\n",
    "results = []\n",
    "for i in range(5):\n",
    "    test_input = torch.rand(1, 3, 224, 224).cuda()\n",
    "    sub_label = f\"[test {i}]\"\n",
    "    results.append(\n",
    "        benchmark.Timer(\n",
    "            stmt=\"resnet18_tvm(test_input)\",\n",
    "            setup=\"from __main__ import resnet18_tvm\",\n",
    "            globals={\"test_input\": test_input},\n",
    "            sub_label=sub_label,\n",
    "            description=\"tuning by meta\",\n",
    "        ).blocked_autorange()\n",
    "    )\n",
    "    results.append(\n",
    "        benchmark.Timer(\n",
    "            stmt=\"resnet18_torch(test_input)\",\n",
    "            setup=\"from __main__ import resnet18_torch\",\n",
    "            globals={\"test_input\": test_input},\n",
    "            sub_label=sub_label,\n",
    "            description=\"tuning by jit\",\n",
    "        ).blocked_autorange()\n",
    "    )\n",
    "\n",
    "compare = benchmark.Compare(results)\n",
    "compare.print()\n",
    "\n",
    "# In author's environment, the average inference time of `resnet18_tvm` is 620.0 us,\n",
    "# while the average inference time of `resnet18_torch` is 980.0 us (PyTorch version is 1.11.0),\n",
    "# showing the speedup of around 38%.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
